"""
内容：输入视频预测数值的模型示例
日期：2020年7月6日
作者：Howie
"""
import numpy as np
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense, Conv2D, Flatten, MaxPooling2D
from keras.utils import plot_model

# 预设
IMG_HEIGHT = 16
IMG_WIDTH = 16
PLT_ROW = 5
PLT_COL = 5

# 超参
N_SAMPLES_TRAIN = 1500
N_SAMPLES_VAL = 300
N_SAMPLES_TEST = 100
EPOCH = 1000
BATCH_SIZE = 32


def hist_plot(hist, model_name):
    """
    # 可视化训练过程
    :param hist: history对象
    :return:
    """
    print(hist.history.keys())
    fig, loss_ax = plt.subplots()
    # 每个训练周期的训练误差与验证误差
    loss_ax.plot(hist.history['loss'], 'y', label='train loss')
    loss_ax.plot(hist.history['val_loss'], 'r', label='val loss')
    # 横轴与纵轴
    loss_ax.set_xlabel('epoch')
    loss_ax.set_ylabel('loss')
    # 标签
    loss_ax.legend(loc='upper left')
    # 保存
    plt.savefig('./logs/History_Demo4_' + model_name + '.pdf')
    # 展示
    plt.show()


def generate_dataset(samples):
    """
    生成数据集:
    准备宽16，高16，像素值分为0和1的视频。输入随机值，重复随机值次数，将视频内值为1的像素拍摄下来
    此处将随机值指定为标签值
    :return:
    """
    ds_x, ds_y = [], []

    for it in range(samples):
        num_pt = np.random.randint(low=0, high=IMG_WIDTH * IMG_HEIGHT)
        img = generate_image(num_pt)

        ds_x.append(img)
        ds_y.append(num_pt)
    return np.array(ds_x), np.array(ds_y).reshape(samples, 1)


def generate_image(points):
    """
    # 生成图片
    :param points:
    :return:
    """
    img = np.zeros((IMG_WIDTH, IMG_HEIGHT))
    pts = np.random.random((points, 2))

    for ipt in pts:
        img[int(ipt[0] * IMG_WIDTH), int(ipt[1] * IMG_HEIGHT)] = 1

    return img.reshape(IMG_WIDTH, IMG_HEIGHT, 1)


def dataset_visualization(sample, label):
    """
    将生成的数据集的一部分进行可视化展示
    :return:
    """
    plt.rcParams["figure.figsize"] = (10, 10)
    fig, ax = plt.subplots(PLT_ROW, PLT_COL)

    for i in range(PLT_ROW * PLT_COL):
        sub_plt = ax[i // PLT_ROW, i % PLT_ROW]
        sub_plt.axis('off')
        sub_plt.imshow(sample[i].reshape(IMG_WIDTH, IMG_HEIGHT))
        sub_plt.set_title('R ' + str(label[i][0]))
    plt.savefig('./logs/Dataset Visualization_Demo4.pdf')
    plt.show()


def dataset_preparation():
    """
    准备数据集
    :return:
    """
    # 训练集
    X_train, Y_train = generate_dataset(N_SAMPLES_TRAIN)
    # 验证集
    X_val, Y_val = generate_dataset(N_SAMPLES_VAL)
    # 测试集
    X_test, Y_test = generate_dataset(N_SAMPLES_TEST)
    # 可视化
    dataset_visualization(X_train, Y_train)

    return X_train, Y_train, X_val, Y_val, X_test, Y_test


def model_building(choice='Multi-Layer Perceptron'):
    """
    # 准备模型
    :return:
    """
    model = Sequential()
    if choice == 'Multi-Layer Perceptron':
        # 多层感应器神经网络模型
        model.add(
            Dense(
                256,
                activation='relu',
                input_dim=IMG_WIDTH *
                IMG_HEIGHT))
        model.add(Dense(256, activation='relu'))
        model.add(Dense(256))
        model.add(Dense(1))
    else:
        # 卷积神经网络
        model.add(
            Conv2D(
                32, (3, 3), activation='relu', input_shape=(
                    IMG_WIDTH, IMG_HEIGHT, 1)))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Conv2D(32, (3, 3), activation='relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Flatten())
        model.add(Dense(256, activation='relu'))
        model.add(Dense(1))
    return model


def model_training():
    # 生成数据集
    X_train, Y_train, X_val, Y_val, X_test, Y_test = dataset_preparation()
    titles = ['Multi-Layer Perceptron',
              'Convolution Neural Network']
    for title in titles:
        # 搭建模型
        model = model_building(choice=title)
        model.compile(optimizer='adam', loss='mse')
        # 训练模型
        X_train_1D = X_train.reshape(N_SAMPLES_TRAIN, IMG_WIDTH * IMG_HEIGHT)
        X_val_1D = X_val.reshape(N_SAMPLES_VAL, IMG_WIDTH * IMG_HEIGHT)
        X_test_1D = X_test.reshape(N_SAMPLES_TEST, IMG_WIDTH * IMG_HEIGHT)
        hist = model.fit(
            X_train if title == 'Convolution Neural Network' else X_train_1D,
            Y_train,
            epochs=EPOCH,
            batch_size=BATCH_SIZE,
            validation_data=(
                X_val if title == 'Convolution Neural Network' else X_val_1D,
                Y_val))
        # 查看训练过程
        hist_plot(hist, title)
        # 评价模型
        score = model.evaluate(
            X_test if title == 'Convolution Neural Network' else X_test_1D,
            Y_test,
            batch_size=BATCH_SIZE)
        print(score)
        # 调用模型
        Y_hat_test = model.predict(
            X_test if title == 'Convolution Neural Network' else X_test_1D,
            batch_size=BATCH_SIZE)
        fig, ax = plt.subplots(PLT_ROW, PLT_COL)

        for i in range(PLT_ROW * PLT_COL):
            sub_plt = ax[i // PLT_ROW, i % PLT_ROW]
            sub_plt.axis('off')
            sub_plt.imshow(X_test[i].reshape(IMG_WIDTH, IMG_HEIGHT))
            sub_plt.set_title('R %d P %.1f' % (Y_test[i][0], Y_hat_test[i][0]))
        plt.savefig('./logs/Predictions of ' + title + '_Demo4.pdf')
        plt.show()


if __name__ == '__main__':
    model_training()
